use crate::ast::*;
use crate::error::{CompilerError, Result};
use super::TypeInference;

impl TypeInference {


    /// Infer types in a declaration
    pub fn infer_declaration(&mut self, declaration: &Declaration) -> Result<()> {
        match declaration {
            Declaration::Function(func) => {
                self.infer_function(func)?;
            }
            Declaration::Variable(var) => {
                self.infer_variable(var)?;
            }
            Declaration::Constant(const_) => {
                self.infer_constant(const_)?;
            }
            Declaration::Type(type_) => {
                self.infer_type_declaration(type_)?;
            }
            Declaration::Struct(struct_) => {
                self.infer_struct(struct_)?;
            }
            Declaration::Enum(_enum) => {
                // For now, enums don't need type inference
            }
            Declaration::Trait(_trait) => {
                // TODO: Infer trait declaration
            }
            Declaration::TraitImpl(trait_impl) => {
                // Infer types for trait implementation
                // Treat trait impl similar to impl blocks
                for method in &trait_impl.methods {
                    // Create new environment for method
                    let mut method_env = self.environment.child();

                    // Add generic parameters to environment
                    for generic in &method.generics {
                        method_env.add_generic(generic.name.clone());
                    }

                    // Infer parameter types
                    for param in &method.parameters {
                        let resolved_type = self.resolve_type(&param.param_type)?;
                        method_env.bind_variable(param.name.clone(), resolved_type);
                    }

                    // Infer return type
                    let return_type = if let Some(return_type) = &method.return_type {
                        return_type.clone()
                    } else {
                        // Infer return type from body
                        if let Some(body) = &method.body {
                            self.infer_block(body, &mut method_env)?
                        } else {
                            Type::Basic(crate::ast::types::BasicType::Void)
                        }
                    };

                    // Create function type
                    let param_types: Vec<Type> = method
                        .parameters
                        .iter()
                        .map(|p| p.param_type.clone())
                        .collect();
                    let function_type = FunctionType {
                        parameter_types: param_types,
                        return_type: Some(Box::new(return_type)),
                        variadic: false,
                        location: method.location,
                    };

                    // Register method with TypeName.methodName format
                    let method_name = format!("{}.{}", trait_impl.for_type, method.name);
                    self.environment.bind_function(method_name, function_type);
                }
            }
            Declaration::Import(import) => {
                self.infer_import(import)?;
            }
            Declaration::Impl(impl_) => {
                self.infer_impl(impl_)?;
            }
        }
        Ok(())
    }


    /// Infer types in a function
    fn infer_function(&mut self, func: &FunctionDecl) -> Result<()> {
        // Create new environment for function
        let mut func_env = self.environment.child();

        // Add generic parameters to environment
        for generic in &func.generics {
            func_env.add_generic(generic.name.clone());
        }

        // Infer parameter types
        for param in &func.parameters {
            // Check if the parameter type is a generic that should be resolved to a struct type
            let resolved_type = self.resolve_type(&param.param_type)?;
            func_env.bind_variable(param.name.clone(), resolved_type);
        }

        let try_depth = self.try_context_stack.len();
        self.push_try_context_from_type(func.return_type.as_ref());
        let return_result = (|| -> Result<Type> {
            if let Some(return_type) = &func.return_type {
                let return_type_clone = return_type.clone();
                // Bind function type early so return statements can use it for context
                let param_types: Vec<Type> = func
                    .parameters
                    .iter()
                    .map(|p| p.param_type.clone())
                    .collect();
                let function_type = FunctionType {
                    parameter_types: param_types.clone(),
                    return_type: Some(Box::new(return_type_clone.clone())),
                    variadic: false,
                    location: func.location,
                };
                func_env.bind_function(func.name.clone(), function_type);

                // Infer function body with expected return type for context-aware inference
                if let Some(body) = &func.body {
                    self.infer_block_with_expected(body, &mut func_env, Some(&return_type_clone))?;
                }

                Ok(return_type_clone)
            } else if let Some(body) = &func.body {
                // Infer return type from body
                self.infer_block(body, &mut func_env)
            } else {
                Ok(Type::Basic(crate::ast::types::BasicType::Void))
            }
        })();
        self.truncate_try_context(try_depth);
        let return_type = return_result?;

        // Create function type (if not already bound)
        let param_types: Vec<Type> = func
            .parameters
            .iter()
            .map(|p| p.param_type.clone())
            .collect();
        let function_type = FunctionType {
            parameter_types: param_types,
            return_type: Some(Box::new(return_type)),
            variadic: false,
            location: func.location,
        };

        // Bind function type (if not already bound)
        if func.return_type.is_none() {
            func_env.bind_function(func.name.clone(), function_type);
        }

        Ok(())
    }


    /// Infer types in a variable declaration
    fn infer_variable(&mut self, var: &VariableDecl) -> Result<()> {
        let var_type = if let Some(declared_type) = &var.var_type {
            declared_type.clone()
        } else if let Some(initializer) = &var.initializer {
            // Infer type from initializer
            let mut env = std::mem::take(&mut self.environment);
            let inferred_type = self.infer_expression(initializer, &mut env)?;
            self.environment = env;
            inferred_type
        } else {
            return Err(CompilerError::type_error(
                var.location.line,
                var.location.column,
                "Variable declaration must have type annotation or initializer",
            ));
        };

        self.environment.bind_variable(var.name.clone(), var_type);
        Ok(())
    }


    /// Infer types in a constant declaration
    fn infer_constant(&mut self, const_: &ConstantDecl) -> Result<()> {
        let const_type = if let Some(declared_type) = &const_.const_type {
            declared_type.clone()
        } else {
            // Infer type from value
            let mut env = std::mem::take(&mut self.environment);
            let inferred_type = self.infer_expression(&const_.value, &mut env)?;
            self.environment = env;
            inferred_type
        };

        self.environment
            .bind_variable(const_.name.clone(), const_type);
        Ok(())
    }


    /// Infer types in a type declaration
    fn infer_type_declaration(&mut self, type_: &TypeDecl) -> Result<()> {
        self.environment
            .bind_type(type_.name.clone(), type_.type_def.clone());
        Ok(())
    }


    /// Infer types in a struct declaration
    fn infer_struct(&mut self, struct_: &StructDecl) -> Result<()> {
        // Create struct type
        let struct_type = Type::Struct(StructType {
            name: struct_.name.clone(),
            type_args: vec![],
            location: struct_.location,
        });

        self.environment
            .bind_type(struct_.name.clone(), struct_type);

        // Infer method types
        for method in &struct_.methods {
            self.infer_function(method)?;
        }

        Ok(())
    }


    /// Infer types in an import declaration
    fn infer_import(&mut self, _import: &ImportDecl) -> Result<()> {
        // TODO: Infer imported types
        Ok(())
    }

    /// Infer types in an implementation declaration
    fn infer_impl(&mut self, impl_: &ImplDecl) -> Result<()> {
        // Infer types for each method in the impl block
        for method in &impl_.methods {
            // Create new environment for method
            let mut method_env = self.environment.child();

            // Add generic parameters to environment
            for generic in &method.generics {
                method_env.add_generic(generic.name.clone());
            }

            // Infer parameter types
            for param in &method.parameters {
                let resolved_type = self.resolve_type(&param.param_type)?;
                method_env.bind_variable(param.name.clone(), resolved_type);
            }

            // Infer return type
            let return_type = if let Some(return_type) = &method.return_type {
                return_type.clone()
            } else {
                // Infer return type from body
                if let Some(body) = &method.body {
                    self.infer_block(body, &mut method_env)?
                } else {
                    Type::Basic(crate::ast::types::BasicType::Void)
                }
            };

            // Create function type
            let param_types: Vec<Type> = method
                .parameters
                .iter()
                .map(|p| p.param_type.clone())
                .collect();
            let function_type = FunctionType {
                parameter_types: param_types,
                return_type: Some(Box::new(return_type)),
                variadic: false,
                location: method.location,
            };

            // Register method with TypeName.methodName format
            let method_name = format!("{}.{}", impl_.type_name, method.name);
            self.environment.bind_function(method_name, function_type);
        }

        Ok(())
    }

}
